import numpy as np 
import pandas as pd 
from pathlib import Path

import pymc as pm
import arviz as az

import cufflinks as cf
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression

from sklearn.metrics import r2_score, mean_squared_error
from ipywidgets import interact, interact_manual, fixed

pd.set_option('display.max_rows', 500)
pd.set_option('use_inf_as_na', True)
cf.set_config_file(offline=True, theme='pearl');
path = Path("data/")

/var/folders/5x/vz4_nh_d6f5ds1c71xs4q_n0qqgwj_/T/ipykernel_73844/2430209213.py:2: FutureWarning:

use_inf_as_na option is deprecated and will be removed in a future version. Convert inf values to NaN before operating instead.

master_df = pd.read_csv(path/'covid_19_data.csv')

recovered_df = (pd.read_csv(path/'time_series_covid_19_recovered.csv')
                .drop(columns=['Lat', 'Long', 'Province/State'])
                .groupby('Country/Region')
                .sum())

deaths_df = (pd.read_csv(path/'time_series_covid_19_deaths.csv')
             .drop(columns=['Lat', 'Long', 'Province/State'])
             .groupby('Country/Region')
             .sum())

confirmed_df = (pd.read_csv(path/'time_series_covid_19_confirmed.csv')
                .drop(columns=['Lat', 'Long', 'Province/State'])
                .groupby('Country/Region')
                .sum())

# sorted_country_list = confirmed_df.sort_values(by=confirmed_df.columns[-1], ascending=False).index.to_list()
sorted_country_list = confirmed_df.index

@interact(country=sorted_country_list, threshold=(1, 1000, 10), fit=True)
def log_lin_visualise(country, fit, threshold=100):
    
    y = confirmed_df.filter(items=[country], axis=0).values.squeeze(0).astype(float)
    y = np.log(y[y > threshold])
    x = np.arange(1, y.shape[0] + 1)
    
    plt.figure(figsize=(10, 5))
    plt.plot(x, y, label='Observed')
    
    if fit:

        lr = LinearRegression()
        lr.fit(X=x.reshape(-1, 1) , y=y)
        α, β =  lr.intercept_, lr.coef_

        y_fitted = lr.predict(X=x.reshape(-1, 1))

        print("Solving linear regression using OLS ... ")
        print(f"* r2_score = {round(r2_score(y, y_fitted), 2)}")
        print(f"* mean_squared_error = {round(mean_squared_error(y, y_fitted), 2)}")

        plt.plot(x, y_fitted, label=f"{round(α, 2)} + {round(β[0], 2)}*x")

    plt.xlabel(f'Days Since {threshold}th Case')
    plt.ylabel('Natural Logarithm of Confirmed Cases')
    plt.legend()
    plt.title(country)
    plt.show()
    plt.close()

interactive(children=(Dropdown(description='country', options=('Afghanistan', 'Albania', 'Algeria', 'Andorra',…

country = 'US'
threshold = 100

y = confirmed_df.filter(items=[country], axis=0).values.squeeze(0)
y = np.log(y[y > threshold])
x = np.arange(1, y.shape[0] + 1)


with pm.Model() as unpooled_model:
    
    
    # priors
    α = pm.Normal(name='α', mu=int(np.log(threshold)), sigma=10)
    β = pm.Normal(name='β')
    
    # error
    σ = pm.HalfNormal(name='σ', sigma=10)
    
    # expected value
    μ = pm.Deterministic(name='μ', var= α + β*x)
    
    # liklihood == 'prior_predictive'
    pm.Normal(name=country, mu=μ, sigma=σ, observed=y)
    
gv = pm.model_to_graphviz(unpooled_model)
gv.format = 'png'
gv.render(filename='unpooled_model.png')
gv

with unpooled_model:
    
    # sampling liklihood
    prior = pm.sample_prior_predictive()
    
    # posterior
    trace = pm.sample()
    
    # predictions == 'posterior_predictive'
    pred = pm.sample_posterior_predictive(trace)

Sampling: [US, α, β, σ]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [α, β, σ]

Output()

Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 1 seconds.
Sampling: [US]

Output()

def custom_from_pymc3(trace, prior, pred):
    unpooled = az.convert_to_inference_data(trace)
    unpooled.extend(pred)
    unpooled.extend(prior)
    return unpooled

# unpooled = az.from_pymc3(trace=trace, prior=prior, posterior_predictive=pred, model=unpooled_model)
unpooled = custom_from_pymc3(trace, prior, pred)
prior_vars = ['α', 'β', 'σ']
unpooled

<xarray.Dataset> Size: 15MB
Dimensions:  (chain: 4, draw: 1000, μ_dim_0: 452)
Coordinates:
  * chain    (chain) int64 32B 0 1 2 3
  * draw     (draw) int64 8kB 0 1 2 3 4 5 6 7 ... 993 994 995 996 997 998 999
  * μ_dim_0  (μ_dim_0) int64 4kB 0 1 2 3 4 5 6 7 ... 445 446 447 448 449 450 451
Data variables:
    α        (chain, draw) float64 32kB 12.45 12.56 12.3 ... 12.5 12.47 12.46
    β        (chain, draw) float64 32kB 0.01308 0.01289 ... 0.0135 0.01359
    σ        (chain, draw) float64 32kB 1.192 1.179 1.285 ... 1.185 1.223 1.274
    μ        (chain, draw, μ_dim_0) float64 14MB 12.47 12.48 ... 18.59 18.61
Attributes:
    created_at:                 2025-01-13T00:15:51.105693+00:00
    arviz_version:              0.20.0
    inference_library:          pymc
    inference_library_version:  5.20.0
    sampling_time:              1.2261970043182373
    tuning_steps:               1000

array([0, 1, 2, 3])

array([  0,   1,   2, ..., 997, 998, 999])

array([  0,   1,   2, ..., 449, 450, 451])

array([[12.45367708, 12.56411582, 12.29812282, ..., 12.46407159,
        12.37800258, 12.37502912],
       [12.50407931, 12.36184365, 12.3443394 , ..., 12.22502794,
        12.74110813, 12.74110813],
       [12.35580532, 12.36260233, 12.63419148, ..., 12.56204529,
        12.58547379, 12.51858901],
       [12.49312726, 12.65033725, 12.56894874, ..., 12.49980907,
        12.46518642, 12.46365533]])

array([[0.0130819 , 0.01288788, 0.01397371, ..., 0.01363426, 0.01355799,
        0.01359591],
       [0.01273469, 0.01332098, 0.0142709 , ..., 0.01428819, 0.01255352,
        0.01255352],
       [0.01367752, 0.013486  , 0.01271762, ..., 0.01284616, 0.01306908,
        0.01304445],
       [0.01306697, 0.01263895, 0.01288684, ..., 0.01338155, 0.01349886,
        0.01358734]])

array([[1.19186607, 1.17914418, 1.28500327, ..., 1.26130331, 1.19639562,
        1.18758237],
       [1.22094158, 1.27568683, 1.20541414, ..., 1.13933684, 1.29352798,
        1.29352798],
       [1.22849687, 1.22408668, 1.26271943, ..., 1.19687879, 1.18001486,
        1.22179074],
       [1.15221658, 1.27939076, 1.23437394, ..., 1.18454095, 1.22256083,
        1.27436151]])

array([[[12.46675898, 12.47984088, 12.49292278, ..., 18.34053246,
         18.35361436, 18.36669626],
        [12.5770037 , 12.58989158, 12.60277946, ..., 18.3636614 ,
         18.37654928, 18.38943716],
        [12.31209653, 12.32607024, 12.34004395, ..., 18.58629282,
         18.60026653, 18.61424024],
        ...,
        [12.47770585, 12.49134011, 12.50497438, ..., 18.59948936,
         18.61312362, 18.62675788],
        [12.39156057, 12.40511856, 12.41867655, ..., 18.47909809,
         18.49265608, 18.50621407],
        [12.38862504, 12.40222095, 12.41581686, ..., 18.49318992,
         18.50678584, 18.52038175]],

       [[12.51681399, 12.52954868, 12.54228337, ..., 18.23468836,
         18.24742305, 18.26015773],
        [12.37516463, 12.38848561, 12.40180659, ..., 18.35628486,
         18.36960584, 18.38292682],
        [12.3586103 , 12.3728812 , 12.3871521 , ..., 18.76624396,
         18.78051486, 18.79478576],
...
        [12.57489145, 12.58773761, 12.60058377, ..., 18.34281739,
         18.35566355, 18.36850971],
        [12.59854287, 12.61161195, 12.62468103, ..., 18.46655913,
         18.4796282 , 18.49269728],
        [12.53163346, 12.54467791, 12.55772236, ..., 18.38859177,
         18.40163622, 18.41468067]],

       [[12.50619424, 12.51926121, 12.53232819, ..., 18.37326591,
         18.38633288, 18.39939986],
        [12.6629762 , 12.67561515, 12.68825411, ..., 18.3378666 ,
         18.35050555, 18.36314451],
        [12.58183559, 12.59472243, 12.60760928, ..., 18.3680283 ,
         18.38091514, 18.39380199],
        ...,
        [12.51319062, 12.52657217, 12.53995372, ..., 18.52150672,
         18.53488827, 18.54826982],
        [12.47868528, 12.49218415, 12.50568301, ..., 18.53967485,
         18.55317371, 18.56667258],
        [12.47724268, 12.49083002, 12.50441736, ..., 18.57796023,
         18.59154758, 18.60513492]]])

PandasIndex(Index([0, 1, 2, 3], dtype='int64', name='chain'))

PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
       ...
       990, 991, 992, 993, 994, 995, 996, 997, 998, 999],
      dtype='int64', name='draw', length=1000))

unpooled.posterior['α'].shape

(4, 1000)

summary_df = az.summary(unpooled)
summary_df.to_csv(f'unpooled_{country}_summary.csv')

az.plot_trace(data=unpooled, var_names=prior_vars);
plt.tight_layout()

az.plot_posterior(data=unpooled, var_names=prior_vars, group='posterior');

az.plot_posterior(data=unpooled, var_names=prior_vars, group='prior');

az.plot_ppc(data=unpooled, group='posterior');

# pm.find_MAP(model=unpooled_model)

mean_priors = az.summary(unpooled).filter(items=prior_vars, axis=0)['mean']
y_fitted = mean_priors['α'] + mean_priors['β']*x

mean_priors

α    12.419
β     0.014
σ     1.236
Name: mean, dtype: float64

ci = 0.95


az.plot_hdi(x=x, y=unpooled.posterior_predictive[country], hdi_prob=ci)#, plot_kwargs={'label': f'Predictions ({ci})'});
# az.plot_hpd(x=x, y=unpooled.posterior['μ'], color='C2', credible_interval=ci)#, plot_kwargs={'label': f'Expectation ({ci})'});
plt.plot(x, y, label='Observed');
plt.plot(x, y_fitted, label=f"{mean_priors['α']} + {mean_priors['β']}*x", color='k');

plt.xlabel(f'Days Since {threshold}th Case')
plt.ylabel('Natural Logarithm of Confirmed Cases')
plt.title(country)
plt.legend()
plt.show()
plt.close()

az.r2_score(y_true=y, y_pred=unpooled.posterior_predictive[country].data.reshape(-1, y.shape[0]))

r2        0.604565
r2_std    0.018961
dtype: float64

# az.plot_dist(unpooled.log_likelihood[country]);
# az.plot_dist(unpooled.prior_predictive[country]);

## Mean of predicted values

y_fitted = unpooled.posterior_predictive[country].data.reshape(-1, y.shape[0]).mean(axis=0)
print(f"r2_score = {round(r2_score(y, y_fitted), 2)}")
print(f"mean_squared_error = {round(mean_squared_error(y, y_fitted), 2)}")

r2_score = 0.67
mean_squared_error = 1.51

## Mean of expected values

y_fitted = unpooled.posterior['μ'].data.reshape(-1, y.shape[0]).mean(axis=0)
print(f"r2_score = {round(r2_score(y, y_fitted), 2)}")
print(f"mean_squared_error = {round(mean_squared_error(y, y_fitted), 2)}")

r2_score = 0.67
mean_squared_error = 1.51

# Data for all counties needs to be in similar dimensions.

cumulative_threshold = 1e4 # to select countries
threshold = 100 # to select start_date

high_confirmed_df = confirmed_df[confirmed_df.iloc[:, -1] > cumulative_threshold]
high_confirmed_df = high_confirmed_df.where(cond=lambda x: x > threshold, other=np.nan) # or -1

# Dropping for now
high_confirmed_df = high_confirmed_df.drop(index=['China', 'Korea, South'])

majority = high_confirmed_df.shape[0]/2

high_confirmed_df.isna().sum().iplot(hline=majority,
                                     xTitle='Date', yTitle=f'Countries below {threshold} cases', 
                                     title='Threshold Cases for Highly Infected Countries');

cross_over_date = (high_confirmed_df.isna().sum() < majority).idxmax()
cross_over_date

'3/23/20'

thresholded_high_confirmed_df = high_confirmed_df.iloc[:, high_confirmed_df.columns.to_list().index(cross_over_date):]
thresholded_high_confirmed_df = thresholded_high_confirmed_df.dropna(how='any', axis=0)
thresholded_high_confirmed_df = thresholded_high_confirmed_df.applymap(np.log)

print(thresholded_high_confirmed_df.shape)

(75, 433)

/var/folders/5x/vz4_nh_d6f5ds1c71xs4q_n0qqgwj_/T/ipykernel_73844/1674293464.py:3: FutureWarning:

DataFrame.applymap has been deprecated. Use DataFrame.map instead.

n_countries = thresholded_high_confirmed_df.shape[0]
date_points = thresholded_high_confirmed_df.shape[1]
idx = np.repeat(a=np.arange(n_countries), repeats=date_points)

x = np.arange(1, date_points + 1)
XX = np.stack([x]*n_countries, axis=0)

YY = thresholded_high_confirmed_df.values

print(XX.shape, YY.shape)

(75, 433) (75, 433)

def split_and_pack(labels, data, axis=2):
    '''
    Helper function to unpack values from a vectorised output
    '''
    splits = np.split(ary=data, axis=axis, indices_or_sections=len(labels))
    d = dict(zip(labels, data))
    return d
    

with pm.Model() as pooled_model:

    
    X = pm.Data(name='X', value=XX)
    Y = pm.Data(name='Y', value=YY.flatten())
    
    # priors
    α = pm.Normal(name='α', mu=int(np.log(threshold)), sigma=10, shape=n_countries)
    β = pm.Normal(name='β', sigma=5, shape=n_countries)
    
    # error
    σ = pm.HalfNormal(name='σ', sigma=10, shape=n_countries)
    
    # expected value
    
    μ = pm.Deterministic(name='μ', var= (α + β*X.T).T.flatten())
    
    # liklihood == 'prior_predictive'
    pm.Normal(name='pooled', mu=μ, sigma=σ[idx], observed=Y)
    
gv = pm.model_to_graphviz(pooled_model)
gv.format = 'png'
gv.render(filename='pooled_model.png')
gv

with pooled_model:
    
    # sampling liklihood
    prior = pm.sample_prior_predictive()
    
    # posterior
    trace = pm.sample()
    
    # predictions == 'posterior_predictive'
    pred = pm.sample_posterior_predictive(trace)

Sampling: [pooled, α, β, σ]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [α, β, σ]

Output()

Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 152 seconds.
Sampling: [pooled]

Output()

# pooled = az.from_pymc3(trace=trace, prior=prior, posterior_predictive=pred, model=pooled_model)
pooled = custom_from_pymc3(trace, prior, pred)
prior_vars = ['α', 'β', 'σ']
pooled

<xarray.Dataset> Size: 1GB
Dimensions:  (chain: 4, draw: 1000, α_dim_0: 75, β_dim_0: 75, σ_dim_0: 75,
              μ_dim_0: 32475)
Coordinates:
  * chain    (chain) int64 32B 0 1 2 3
  * draw     (draw) int64 8kB 0 1 2 3 4 5 6 7 ... 993 994 995 996 997 998 999
  * α_dim_0  (α_dim_0) int64 600B 0 1 2 3 4 5 6 7 8 ... 67 68 69 70 71 72 73 74
  * β_dim_0  (β_dim_0) int64 600B 0 1 2 3 4 5 6 7 8 ... 67 68 69 70 71 72 73 74
  * σ_dim_0  (σ_dim_0) int64 600B 0 1 2 3 4 5 6 7 8 ... 67 68 69 70 71 72 73 74
  * μ_dim_0  (μ_dim_0) int64 260kB 0 1 2 3 4 5 ... 32470 32471 32472 32473 32474
Data variables:
    α        (chain, draw, α_dim_0) float64 2MB 6.317 8.314 ... 11.39 5.156
    β        (chain, draw, β_dim_0) float64 2MB 0.01483 0.0102 ... 0.01642
    σ        (chain, draw, σ_dim_0) float64 2MB 0.4209 0.6319 ... 0.3675 0.433
    μ        (chain, draw, μ_dim_0) float64 1GB 6.332 6.347 ... 12.25 12.27
Attributes:
    created_at:                 2025-01-13T00:18:37.416550+00:00
    arviz_version:              0.20.0
    inference_library:          pymc
    inference_library_version:  5.20.0
    sampling_time:              151.56448006629944
    tuning_steps:               1000

array([0, 1, 2, 3])

array([  0,   1,   2, ..., 997, 998, 999])

array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
       18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
       36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
       54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
       72, 73, 74])

array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
       18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
       36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
       54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
       72, 73, 74])

array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
       18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
       36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
       54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
       72, 73, 74])

array([    0,     1,     2, ..., 32472, 32473, 32474])

array([[[ 6.31691393,  8.31387285,  5.92331806, ...,  9.0766695 ,
         11.37603529,  5.14027915],
        [ 6.31478642,  8.4770691 ,  6.02745653, ...,  9.05059758,
         11.36252844,  5.22595593],
        [ 6.25013637,  8.34045168,  5.93983991, ...,  9.02836941,
         11.34346115,  5.14309506],
        ...,
        [ 6.32859896,  8.42679732,  5.96670203, ...,  9.03919413,
         11.40936324,  5.18713252],
        [ 6.2636508 ,  8.45058991,  5.99526093, ...,  9.11585365,
         11.3097885 ,  5.15674207],
        [ 6.32153351,  8.32868062,  5.96550581, ...,  9.12325828,
         11.37692161,  5.16359064]],

       [[ 6.28158279,  8.44035345,  5.96487312, ...,  9.07380898,
         11.30395793,  5.16959017],
        [ 6.27611919,  8.36365347,  5.98314925, ...,  9.11493728,
         11.38810529,  5.16143501],
        [ 6.30336029,  8.38778233,  5.95773965, ...,  9.0620908 ,
         11.37856256,  5.26847899],
...
        [ 6.31333001,  8.35246288,  5.9565986 , ...,  9.07977628,
         11.4308989 ,  5.08463483],
        [ 6.31742569,  8.49742904,  5.97592816, ...,  9.00543723,
         11.33732651,  5.19795037],
        [ 6.28995354,  8.49640698,  5.95205771, ...,  9.03720172,
         11.38341967,  5.18499343]],

       [[ 6.30867786,  8.27152328,  5.96265849, ...,  9.12280113,
         11.36707229,  5.20449769],
        [ 6.28345822,  8.52808958,  6.00085482, ...,  9.05079495,
         11.36124368,  5.09259903],
        [ 6.27009362,  8.29963406,  5.92168439, ...,  9.03214697,
         11.3266815 ,  5.22146027],
        ...,
        [ 6.2660602 ,  8.42009898,  5.95104725, ...,  9.13351842,
         11.30902959,  5.16350511],
        [ 6.28619862,  8.47101513,  5.94508924, ...,  9.07234032,
         11.35665812,  5.24463571],
        [ 6.26792546,  8.46476489,  5.99872919, ...,  9.01718872,
         11.3921894 ,  5.15618754]]])

array([[[0.01483014, 0.01019741, 0.00958662, ..., 0.01110848,
         0.01061829, 0.01653686],
        [0.01521238, 0.0096571 , 0.00926914, ..., 0.01111488,
         0.01063738, 0.01611085],
        [0.01532593, 0.01021439, 0.0094737 , ..., 0.01139287,
         0.01078456, 0.01652823],
        ...,
        [0.01484557, 0.00975152, 0.00933508, ..., 0.0112247 ,
         0.01054916, 0.01617431],
        [0.01534808, 0.00955862, 0.00946186, ..., 0.01078144,
         0.01083359, 0.01656007],
        [0.01513039, 0.01001957, 0.00943598, ..., 0.01099716,
         0.01055188, 0.01629988]],

       [[0.01502454, 0.00982832, 0.00940193, ..., 0.01121118,
         0.01085761, 0.01660593],
        [0.01524443, 0.00973135, 0.00947047, ..., 0.01108461,
         0.01053511, 0.01616728],
        [0.01500792, 0.01003492, 0.00940566, ..., 0.01094903,
         0.01057085, 0.01621733],
...
        [0.01503321, 0.01006891, 0.00950639, ..., 0.01109861,
         0.01043533, 0.01670329],
        [0.01492801, 0.0093481 , 0.00942344, ..., 0.01137901,
         0.01077279, 0.01610114],
        [0.01502546, 0.00922779, 0.00957547, ..., 0.011086  ,
         0.01067136, 0.01623199]],

       [[0.01496874, 0.0100849 , 0.00954705, ..., 0.01090699,
         0.01073127, 0.01631352],
        [0.0151532 , 0.009381  , 0.00938303, ..., 0.01130578,
         0.01073325, 0.01653878],
        [0.01522175, 0.01004937, 0.0097071 , ..., 0.01126542,
         0.0107748 , 0.01628613],
        ...,
        [0.01529598, 0.00965334, 0.00945674, ..., 0.01098151,
         0.01077303, 0.01631499],
        [0.015145  , 0.00976846, 0.00956045, ..., 0.01106834,
         0.01073148, 0.01617712],
        [0.0152553 , 0.00972533, 0.00933668, ..., 0.0113437 ,
         0.01047754, 0.01641982]]])

array([[[0.42085581, 0.63188942, 0.29553239, ..., 0.62747715,
         0.3866076 , 0.41399691],
        [0.44367956, 0.6296604 , 0.33083922, ..., 0.6517164 ,
         0.35233701, 0.48804334],
        [0.40625335, 0.61191057, 0.29794705, ..., 0.63046528,
         0.37855063, 0.42848316],
        ...,
        [0.42673834, 0.62606646, 0.30397437, ..., 0.60249933,
         0.36085106, 0.44031694],
        [0.41932362, 0.64254697, 0.32187976, ..., 0.68322315,
         0.3719321 , 0.46876678],
        [0.43500941, 0.66798667, 0.30716589, ..., 0.64317168,
         0.3735603 , 0.45631967]],

       [[0.43867033, 0.62585171, 0.29863717, ..., 0.66408179,
         0.39434263, 0.41813713],
        [0.41548846, 0.66602858, 0.32453349, ..., 0.60354785,
         0.34661915, 0.48425433],
        [0.39178065, 0.62280515, 0.31388449, ..., 0.6394162 ,
         0.36481801, 0.45776884],
...
        [0.41810022, 0.60141541, 0.29653835, ..., 0.61989498,
         0.37213024, 0.45398653],
        [0.4249729 , 0.66753299, 0.31717593, ..., 0.65476915,
         0.35756621, 0.44294016],
        [0.42442647, 0.62787021, 0.32705048, ..., 0.59599789,
         0.36971702, 0.46530816]],

       [[0.41076808, 0.57303194, 0.31179708, ..., 0.62864231,
         0.39938297, 0.45001163],
        [0.41695329, 0.65000712, 0.32243357, ..., 0.63548248,
         0.34661613, 0.47313213],
        [0.43905634, 0.60699055, 0.30023469, ..., 0.64668795,
         0.39282731, 0.43917474],
        ...,
        [0.45125653, 0.61861198, 0.30478863, ..., 0.65764537,
         0.39057909, 0.46928031],
        [0.4486783 , 0.65689595, 0.32825562, ..., 0.64921546,
         0.38070076, 0.43125609],
        [0.44911618, 0.66109146, 0.32617593, ..., 0.65022334,
         0.36745918, 0.43296165]]])

# pm.find_MAP(model=pooled_model)

pooled.posterior['α'].shape # chains x samples x n_countries

(4, 1000, 75)

summary_df = az.summary(pooled) 
summary_df.to_csv('pooled_summary.csv')

countries = thresholded_high_confirmed_df.index.to_list()
means = {var : split_and_pack(data=summary_df.filter(like=var, axis=0)['mean'], labels=countries, axis=0) for var in prior_vars }
means_df = pd.DataFrame(means)

means_df

/Users/22337766/anaconda3/envs/pymc_env/lib/python3.12/site-packages/numpy/core/fromnumeric.py:59: FutureWarning:

'Series.swapaxes' is deprecated and will be removed in a future version. Please use 'Series.transpose' instead.

az.plot_trace(data=pooled, var_names=prior_vars);

# # Plotting posterior will give a combined view of all countries
# az.plot_posterior(data=pooled, var_names=prior_vars, group='posterior');

# az.plot_posterior(data=pooled, var_names=prior_vars, group='prior');

az.plot_forest(data=pooled, var_names=prior_vars, combined=True, hdi_prob=0.99);
plt.tight_layout()

# posterior_predictive = split_and_pack(data=pooled.posterior_predictive['pooled'].data, labels=countries, axis=2) # not working??
posterior_predictive = dict(zip(countries, np.split(ary=pooled.posterior_predictive['pooled'].data, axis=2, indices_or_sections=len(countries))))

X_countrywise = split_and_pack(data=pooled.constant_data['X'].data, labels=countries, axis=0)

# Y_countrywise = split_and_pack(data=pooled.constant_data['Y'].data, labels=countries, axis=0)
Y_countrywise = dict(zip(countries, np.split(ary=pooled.constant_data['Y'].data, axis=0, indices_or_sections=len(countries))))

@interact(country=countries)
def plot_countrywise_posterior(country):
    
    ci = 0.95
    
    plt.figure(figsize=(15, 10))

    az.plot_hdi(x=x, y=posterior_predictive[country], hdi_prob=ci);
    plt.plot(X_countrywise[country], Y_countrywise[country], label='Observed');
    
    y_fitted = means_df['α'][country] + means_df['β'][country]*X_countrywise[country]
    
    plt.plot(X_countrywise[country], y_fitted, label=f"{means_df['α'][country]} + {means_df['β'][country]}*x", color='k');

    plt.xlabel(f'Days Since {threshold}th Case')
    plt.ylabel('Natural Logarithm of Confirmed Cases')
    plt.title(country)
    plt.legend()
    plt.show()
    plt.close()
    
    bayesian_r2 = az.r2_score(y_true=Y_countrywise[country], y_pred=posterior_predictive[country].reshape(-1, Y_countrywise[country].shape[0]))
    y_fitted_bayesian = posterior_predictive[country].reshape(-1, Y_countrywise[country].shape[0]).mean(axis=0)
    print(f"Bayesian r2_score = {round(bayesian_r2['r2'], 2)}")
    print(f"Bayesian r2_score std = {round(bayesian_r2['r2_std'], 2)}")
    print('\n')
    print(f"Point(mean) r2_score = {round(r2_score(Y_countrywise[country], y_fitted_bayesian), 2)}")
    print(f"Point(mean) mean_squared_error = {round(mean_squared_error(Y_countrywise[country], y_fitted_bayesian), 2)}")

interactive(children=(Dropdown(description='country', options=('Albania', 'Algeria', 'Andorra', 'Argentina', '…

x = np.arange(1, date_points + 1)
XX = np.stack([x]*n_countries, axis=0)

YY = thresholded_high_confirmed_df.values

print(XX.shape, YY.shape)

(75, 433) (75, 433)

# non-vectorised implementation
'''
All three parameters (mu, sd, observed) have to be of the same shape / broadcasteble
'''

with pm.Model() as hierarchical_model:

    X = pm.Data(name='X', value=XX)
    Y = pm.Data(name='Y', value=YY) # not using it since we will iterate over the dataframe
    
    # hyper-priors
    α_μ = pm.Normal(name='α_μ', mu=int(np.log(threshold)), sigma=10)
    α_σ = pm.HalfNormal(name='α_σ', sigma=10)
    
    β_μ = pm.Normal(name='β_μ', sigma=10)
    β_σ = pm.HalfNormal(name='β_σ', sigma=10)
    
    # priors
    α = pm.Normal(name='α', mu=α_μ, sigma=α_σ, shape=n_countries)
    β = pm.Normal(name='β', mu=β_μ, sigma=β_σ, shape=n_countries)
    
    # error
    σ = pm.HalfNormal(name='σ', sigma=10, shape=n_countries)
    
    # liklihood
    for i, (index, row) in enumerate(thresholded_high_confirmed_df.iterrows()):
        
        μ = pm.Deterministic(name=f'μ_{index}', var= α[i] + β[i]*X[i])
        zeros = np.zeros(shape=(YY.shape[1], ))
        pm.Normal(name=index, mu=μ-Y[i], sigma=σ[i], observed=zeros)  # or observed=row.values
                    
gv = pm.model_to_graphviz(hierarchical_model)
gv.format = 'png'
gv.render(filename='hierarchical_model.png')
gv

with hierarchical_model:
    
    # sampling liklihood
    prior = pm.sample_prior_predictive()
    
    # posterior
    trace = pm.sample()
    
    # predictions == 'posterior_predictive'
    pred = pm.sample_posterior_predictive(trace)

Sampling: [Albania, Algeria, Andorra, Argentina, Armenia, Australia, Austria, Bahrain, Belgium, Bosnia and Herzegovina, Brazil, Bulgaria, Canada, Chile, Colombia, Costa Rica, Croatia, Cyprus, Czechia, Denmark, Dominican Republic, Ecuador, Egypt, Estonia, Finland, France, Germany, Greece, Hungary, India, Indonesia, Iran, Iraq, Ireland, Israel, Italy, Japan, Jordan, Kuwait, Latvia, Lebanon, Lithuania, Luxembourg, Malaysia, Malta, Mexico, Moldova, Morocco, Netherlands, North Macedonia, Norway, Pakistan, Panama, Peru, Philippines, Poland, Portugal, Qatar, Romania, Russia, Saudi Arabia, Serbia, Singapore, Slovakia, Slovenia, South Africa, Spain, Sweden, Switzerland, Thailand, Turkey, US, United Arab Emirates, United Kingdom, Uruguay, α, α_μ, α_σ, β, β_μ, β_σ, σ]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [α_μ, α_σ, β_μ, β_σ, α, β, σ]

Output()

Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 181 seconds.
Sampling: [Albania, Algeria, Andorra, Argentina, Armenia, Australia, Austria, Bahrain, Belgium, Bosnia and Herzegovina, Brazil, Bulgaria, Canada, Chile, Colombia, Costa Rica, Croatia, Cyprus, Czechia, Denmark, Dominican Republic, Ecuador, Egypt, Estonia, Finland, France, Germany, Greece, Hungary, India, Indonesia, Iran, Iraq, Ireland, Israel, Italy, Japan, Jordan, Kuwait, Latvia, Lebanon, Lithuania, Luxembourg, Malaysia, Malta, Mexico, Moldova, Morocco, Netherlands, North Macedonia, Norway, Pakistan, Panama, Peru, Philippines, Poland, Portugal, Qatar, Romania, Russia, Saudi Arabia, Serbia, Singapore, Slovakia, Slovenia, South Africa, Spain, Sweden, Switzerland, Thailand, Turkey, US, United Arab Emirates, United Kingdom, Uruguay]

Output()

# hierarchical = az.from_pymc3(trace=trace, prior=prior, posterior_predictive=pred, model=hierarchical_model)
hierarchical = custom_from_pymc3(trace, prior, pred)
prior_vars = ['α', 'β', 'σ']
hierarchical

<xarray.Dataset> Size: 1GB
Dimensions:                         (chain: 4, draw: 1000, α_dim_0: 75,
                                     β_dim_0: 75, σ_dim_0: 75,
                                     μ_Albania_dim_0: 433,
                                     μ_Algeria_dim_0: 433,
                                     μ_Andorra_dim_0: 433,
                                     μ_Argentina_dim_0: 433,
                                     ...
                                     μ_Switzerland_dim_0: 433,
                                     μ_Thailand_dim_0: 433,
                                     μ_Turkey_dim_0: 433, μ_US_dim_0: 433,
                                     μ_United Arab Emirates_dim_0: 433,
                                     μ_United Kingdom_dim_0: 433,
                                     μ_Uruguay_dim_0: 433)
Coordinates: (12/80)
  * chain                           (chain) int64 32B 0 1 2 3
  * draw                            (draw) int64 8kB 0 1 2 3 ... 996 997 998 999
  * α_dim_0                         (α_dim_0) int64 600B 0 1 2 3 ... 71 72 73 74
  * β_dim_0                         (β_dim_0) int64 600B 0 1 2 3 ... 71 72 73 74
  * σ_dim_0                         (σ_dim_0) int64 600B 0 1 2 3 ... 71 72 73 74
  * μ_Albania_dim_0                 (μ_Albania_dim_0) int64 3kB 0 1 ... 431 432
    ...                              ...
  * μ_Thailand_dim_0                (μ_Thailand_dim_0) int64 3kB 0 1 ... 431 432
  * μ_Turkey_dim_0                  (μ_Turkey_dim_0) int64 3kB 0 1 2 ... 431 432
  * μ_US_dim_0                      (μ_US_dim_0) int64 3kB 0 1 2 ... 430 431 432
  * μ_United Arab Emirates_dim_0    (μ_United Arab Emirates_dim_0) int64 3kB ...
  * μ_United Kingdom_dim_0          (μ_United Kingdom_dim_0) int64 3kB 0 ... 432
  * μ_Uruguay_dim_0                 (μ_Uruguay_dim_0) int64 3kB 0 1 ... 431 432
Data variables: (12/82)
    α_μ                             (chain, draw) float64 32kB 8.587 ... 8.547
    β_μ                             (chain, draw) float64 32kB 0.01199 ... 0....
    α                               (chain, draw, α_dim_0) float64 2MB 6.308 ...
    β                               (chain, draw, β_dim_0) float64 2MB 0.0150...
    α_σ                             (chain, draw) float64 32kB 1.945 ... 1.908
    β_σ                             (chain, draw) float64 32kB 0.004494 ... 0...
    ...                              ...
    μ_Thailand                      (chain, draw, μ_Thailand_dim_0) float64 14MB ...
    μ_Turkey                        (chain, draw, μ_Turkey_dim_0) float64 14MB ...
    μ_US                            (chain, draw, μ_US_dim_0) float64 14MB 13...
    μ_United Arab Emirates          (chain, draw, μ_United Arab Emirates_dim_0) float64 14MB ...
    μ_United Kingdom                (chain, draw, μ_United Kingdom_dim_0) float64 14MB ...
    μ_Uruguay                       (chain, draw, μ_Uruguay_dim_0) float64 14MB ...
Attributes:
    created_at:                 2025-01-13T00:23:58.585027+00:00
    arviz_version:              0.20.0
    inference_library:          pymc
    inference_library_version:  5.20.0
    sampling_time:              180.69681096076965
    tuning_steps:               1000

array([0, 1, 2, 3])

array([  0,   1,   2, ..., 997, 998, 999])

array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
       18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
       36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
       54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
       72, 73, 74])

array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
       18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
       36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
       54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
       72, 73, 74])

array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
       18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
       36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53,
       54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71,
       72, 73, 74])

array([  0,   1,   2, ..., 430, 431, 432])

array([  0,   1,   2, ..., 430, 431, 432])

array([  0,   1,   2, ..., 430, 431, 432])

array([  0,   1,   2, ..., 430, 431, 432])

# pm.find_MAP(model=hierarchical_model)

hierarchical.posterior['α'].shape # chains x samples x n_countries

(4, 1000, 75)

summary_df = az.summary(hierarchical) 
summary_df.to_csv('hierarchical_summary.csv')

countries = thresholded_high_confirmed_df.index.to_list()
means = {var : split_and_pack(data=summary_df.filter(like=f"{var}[", axis=0)['mean'], labels=countries, axis=0) for var in prior_vars }
means_df = pd.DataFrame(means)

means_df

/Users/22337766/anaconda3/envs/pymc_env/lib/python3.12/site-packages/numpy/core/fromnumeric.py:59: FutureWarning:

'Series.swapaxes' is deprecated and will be removed in a future version. Please use 'Series.transpose' instead.

az.plot_trace(data=hierarchical, var_names=prior_vars);

# # Plotting posterior will give a combined view of all countries
# az.plot_posterior(data=hierarchical, var_names=prior_vars, group='posterior');

# az.plot_posterior(data=hierarchical, var_names=prior_vars, group='prior');

az.plot_forest(data=hierarchical, var_names=prior_vars, combined=True, hdi_prob=0.99);

X_countrywise = split_and_pack(data=hierarchical.constant_data['X'].data, labels=countries, axis=0)
Y_countrywise = split_and_pack(data=hierarchical.constant_data['Y'].data, labels=countries, axis=0)

@interact(country=countries)
def plot_countrywise_posterior(country):
    
    ci = 0.95
    
    plt.figure(figsize=(15, 10))

    az.plot_hdi(x=X_countrywise[country], y=hierarchical.posterior_predictive[country], hdi_prob=ci);
    plt.plot(X_countrywise[country], Y_countrywise[country], label='Observed');
    
    y_fitted = means_df['α'][country] + means_df['β'][country]*X_countrywise[country]
    
    plt.plot(X_countrywise[country], y_fitted, label=f"{means_df['α'][country]} + {means_df['β'][country]}*x", color='k');

    plt.xlabel(f'Days Since {threshold}th Case')
    plt.ylabel('Natural Logarithm of Confirmed Cases')
    plt.title(country)
    plt.legend()
    plt.show()
    plt.close()
    
    bayesian_r2 = az.r2_score(y_true=Y_countrywise[country], y_pred=hierarchical.posterior_predictive[country].data.reshape(-1, Y_countrywise[country].shape[0]))
    y_fitted_bayesian = hierarchical.posterior_predictive[country].data.reshape(-1, Y_countrywise[country].shape[0]).mean(axis=0)
    print(f"Bayesian r2_score = {round(bayesian_r2['r2'], 2)}")
    print(f"Bayesian r2_score std = {round(bayesian_r2['r2_std'], 2)}")
    print('\n')
    print(f"Point(mean) r2_score = {round(r2_score(Y_countrywise[country], y_fitted_bayesian), 2)}")
    print(f"Point(mean) mean_squared_error = {round(mean_squared_error(Y_countrywise[country], y_fitted_bayesian), 2)}")

interactive(children=(Dropdown(description='country', options=('Albania', 'Algeria', 'Andorra', 'Argentina', '…

datasets = {'unpooled' : unpooled,
           'pooled' : pooled,
           'hierarchical': hierarchical}

az.plot_forest(data=list(datasets.values()), model_names=list(datasets.keys()), var_names=prior_vars, hdi_prob=0.99, combined=True);
plt.savefig('comparison.png')

	α	β	σ
Albania	6.287	0.015	0.427
Algeria	8.401	0.010	0.629
Andorra	5.971	0.009	0.312
Argentina	8.949	0.017	1.014
Armenia	8.171	0.012	0.753
Australia	8.880	0.004	0.349
Austria	8.859	0.012	0.397
Bahrain	8.473	0.010	0.732
Belgium	10.071	0.010	0.397
Bosnia and Herzegovina	7.132	0.014	0.528
Brazil	11.888	0.013	1.048
Bulgaria	6.838	0.017	0.492
Canada	10.286	0.010	0.381
Chile	10.508	0.010	0.870
Colombia	9.387	0.016	1.039
Costa Rica	6.876	0.016	0.883
Croatia	6.780	0.016	0.500
Cyprus	5.606	0.014	0.408
Czechia	8.027	0.017	0.528
Denmark	8.405	0.011	0.311
Dominican Republic	9.012	0.010	0.699
Ecuador	9.749	0.009	0.558
Egypt	9.196	0.009	0.839
Estonia	6.258	0.013	0.428
Finland	7.831	0.009	0.248
France	11.198	0.011	0.345
Germany	11.223	0.010	0.295
Greece	7.010	0.015	0.370
Hungary	6.696	0.018	0.577
India	10.988	0.017	1.237
Indonesia	9.145	0.015	0.632
Iran	11.249	0.009	0.213
Iraq	8.519	0.016	1.058
Ireland	9.094	0.009	0.353
Israel	9.318	0.012	0.476
Italy	11.460	0.010	0.331
Japan	8.837	0.012	0.335
Jordan	5.118	0.023	0.793
Kuwait	8.817	0.011	0.920
Latvia	5.522	0.016	0.528
Lebanon	6.179	0.019	0.566
Lithuania	5.969	0.017	0.604
Luxembourg	7.657	0.009	0.285
Malaysia	7.690	0.013	0.398
Malta	5.471	0.013	0.309
Mexico	10.194	0.013	0.990
Moldova	8.047	0.012	0.630
Morocco	8.073	0.015	0.749
Netherlands	9.744	0.012	0.350
North Macedonia	7.061	0.013	0.492
Norway	8.302	0.008	0.206
Pakistan	10.276	0.010	0.838
Panama	8.844	0.012	0.657
Peru	10.651	0.011	1.015
Philippines	9.177	0.013	0.720
Poland	8.689	0.017	0.506
Portugal	9.402	0.011	0.338
Qatar	9.804	0.007	0.871
Romania	8.899	0.014	0.488
Russia	11.335	0.012	0.975
Saudi Arabia	10.465	0.008	0.960
Serbia	8.083	0.014	0.470
Singapore	9.610	0.005	0.745
Slovakia	5.997	0.019	0.605
Slovenia	6.136	0.017	0.598
South Africa	9.695	0.014	1.122
Spain	11.701	0.009	0.254
Sweden	9.521	0.011	0.350
Switzerland	9.520	0.010	0.396
Thailand	7.049	0.008	0.515
Turkey	10.703	0.012	0.498
US	13.578	0.010	0.453
United Arab Emirates	9.084	0.011	0.641
United Kingdom	11.350	0.011	0.369
Uruguay	5.186	0.016	0.454

	α	β	σ
Albania	6.289	0.015	0.427
Algeria	8.398	0.010	0.629
Andorra	5.971	0.009	0.312
Argentina	8.960	0.017	1.015
Armenia	8.170	0.012	0.752
Australia	8.876	0.004	0.349
Austria	8.859	0.012	0.397
Bahrain	8.471	0.010	0.732
Belgium	10.069	0.010	0.396
Bosnia and Herzegovina	7.134	0.014	0.529
Brazil	11.883	0.013	1.048
Bulgaria	6.841	0.017	0.492
Canada	10.284	0.010	0.380
Chile	10.500	0.010	0.869
Colombia	9.397	0.016	1.039
Costa Rica	6.888	0.016	0.883
Croatia	6.783	0.016	0.500
Cyprus	5.608	0.014	0.408
Czechia	8.033	0.017	0.529
Denmark	8.404	0.011	0.311
Dominican Republic	9.008	0.010	0.699
Ecuador	9.745	0.009	0.559
Egypt	9.191	0.009	0.839
Estonia	6.260	0.013	0.428
Finland	7.831	0.009	0.248
France	11.197	0.011	0.345
Germany	11.222	0.010	0.295
Greece	7.013	0.015	0.370
Hungary	6.702	0.018	0.577
India	11.000	0.017	1.237
Indonesia	9.147	0.015	0.632
Iran	11.248	0.009	0.213
Iraq	8.530	0.016	1.058
Ireland	9.093	0.009	0.353
Israel	9.317	0.012	0.476
Italy	11.458	0.010	0.331
Japan	8.837	0.012	0.335
Jordan	5.141	0.023	0.794
Kuwait	8.815	0.011	0.920
Latvia	5.528	0.016	0.528
Lebanon	6.186	0.019	0.566
Lithuania	5.976	0.017	0.603
Luxembourg	7.656	0.009	0.285
Malaysia	7.691	0.013	0.398
Malta	5.473	0.013	0.309
Mexico	10.190	0.013	0.990
Moldova	8.048	0.012	0.629
Morocco	8.079	0.015	0.749
Netherlands	9.743	0.012	0.350
North Macedonia	7.062	0.013	0.493
Norway	8.302	0.008	0.206
Pakistan	10.268	0.010	0.838
Panama	8.843	0.012	0.657
Peru	10.639	0.011	1.014
Philippines	9.178	0.013	0.720
Poland	8.693	0.017	0.506
Portugal	9.401	0.011	0.338
Qatar	9.791	0.007	0.871
Romania	8.901	0.014	0.487
Russia	11.329	0.012	0.975
Saudi Arabia	10.447	0.008	0.960
Serbia	8.085	0.014	0.470
Singapore	9.598	0.005	0.745
Slovakia	6.005	0.019	0.605
Slovenia	6.143	0.017	0.598
South Africa	9.697	0.014	1.121
Spain	11.700	0.009	0.254
Sweden	9.520	0.011	0.350
Switzerland	9.518	0.010	0.396
Thailand	7.047	0.008	0.514
Turkey	10.702	0.012	0.498
US	13.575	0.010	0.453
United Arab Emirates	9.081	0.011	0.641
United Kingdom	11.349	0.011	0.369
Uruguay	5.190	0.016	0.454

Objectives:¶

Rationale¶

Observation¶

Model¶

Unpooled Model¶

Finding appropriate start date for all data¶

Preparing data¶

Pooled Model¶

Hierarchical Model¶

Comparison¶

To-Do:¶